# -*- coding: utf-8 -*-
"""
Created on Tue Dec 21 11:06:28 2021

@author: xiong
"""
import numpy as np
import matplotlib.pylab as plt
import collections
import pandas as pd
import csv
import os
import re
from matplotlib.font_manager import FontProperties
from matplotlib.ticker import MultipleLocator
from matplotlib.ticker import FormatStrFormatter

font_set = FontProperties(fname=r'C:\windows\fonts\simsun.ttc',size=12)

def read_csv(file_path):
    with open(file_path,encoding="utf-8") as f:
        csv_reader = csv.reader(f)
        tb = []
        E = []
        type_dict = collections.defaultdict(list)
        for ri,row in enumerate(csv_reader):
            if ri == 0:
                continue
            tb.append(float(row[0]))
            E.append(float(row[1]))

            row_total = sum([float(i) for i in row[2:]])
            for i,num in enumerate(row[2:]):
                type_dict[i+1].append(float(num)/row_total)
    
    
    tb = np.array(tb)
    E = np.array(E)
    for i in type_dict[i]:
        type_dict[i] = np.array(type_dict[i])

    tb = tb - 0.5/2   #调整到各"自中心点
    E = E - (5.893*1.8099*0.5)/2
    
    
    if min(E)>-5.893*1.8099/2 *5:
        a = -5.893*1.8099/2 *5 
    elif min(E)>-5.893*1.8099/2 *25:
        a = -5.893*1.8099/2 *25
    else:
        a = -5.893*1.8099/2 *35
    if max(E) <= 5.893*1.8099/2 *5 :
        if min(E)>-5.893*1.8099/2 *25:
            a = -5.893*1.8099/2 *25
        else:
            a = -5.893*1.8099/2 *35
        b = 5.893*1.8099/2 *5
        sign_loc = 'lower right'
    elif max(E) < 5.893*1.8099/2 *35:
        b = 5.893*1.8099/2 *35
        sign_loc = 'upper right'
    elif  max(E) < 5.893*1.8099/2 *75:
        b = 5.893*1.8099/2 *75
        sign_loc = 'upper right'
    elif max(E) < 5.893*1.8099/2 *175:
        b = 5.893*1.8099/2 *175
        sign_loc = 'upper right'
    return tb,E,type_dict, (a,b),round((b-a)/(5.893*1.8099/2)/10),sign_loc
    a = """   
    if max(E)-min(E) <= 5.893*1.8099*20:
        a = min(E) - 5.893*1.8099*10+(max(E)-min(E))/2    
        b = max(E) + 5.893*1.8099*10-(max(E)-min(E))/2
        a = a - a%(5.893*1.8099/2) - (5.893*1.8099/2)
        b = b - b%(5.893*1.8099/2) + (5.893*1.8099/2)
        return tb,E,type_dict, (a,b),2
    else:
        slim = int((max(E)-min(E))/(5.893*1.8099*10))
        return tb,E,type_dict, (min(E)-(5.893*1.8099*10),max(E)+(5.893*1.8099*10)),slim+2"""
     
             
            
class LatticeImage:
    def __init__(self,csv_path,title):
        
        self.color = ["y",'r','darkorchid','g','b','fuchsia','limegreen','lightseagreen','chocolate','deepskyblue','y','#990033','#FF9966','#996699','#FF99CC','#999900','#50616d']

        self.markerstyle = '.'
        
        self.tb,self.E,self.type_dict,y_range,yy ,self.sign_loc= read_csv(csv_path)
        
        self.type_num = max(self.type_dict.keys())
        
        self.TB_max,self.TB_min,self.E_max,self.E_min = max(self.tb),min(self.tb),max(self.E),min(self.E)
        
        self.fig= plt.subplots(figsize=(12,yy*2))
        self.ax = plt.gca()
        self.ax.xaxis.set_ticks_position('bottom')
        self.ax.spines['bottom'].set_position(('data', 0))
        self.ax.yaxis.set_ticks_position('left')
        self.ax.spines['left'].set_position(('data', 0))
        self.ax.spines['right'].set_color('none')
        self.ax.spines['top'].set_color('none')
        self.ax.plot(np.linspace(0, 30, 1000000), 5.893*1.8099*np.linspace(0, 30, 1000000), '--k',label='分区界限')
        self.ax.plot(np.linspace(0, 30, 1000000), -5.893*1.8099*np.linspace(0, 30, 1000000), '--k')        
        plt.axhline(0, color='k')#

        self.ax.xaxis.grid(True, which='minor')
        self.ax.xaxis.grid(True)
        self.ax.yaxis.grid(True, which='minor')
        self.ax.yaxis.grid(True)
        self.ax.set_xlabel('TB (℃)',fontproperties =font_set)
        self.ax.set_ylabel('E',fontproperties =font_set)
        self.ax.set_title(title, fontproperties = font_set)
        self.ax.xaxis.set_major_locator(MultipleLocator(5))
        self.ax.yaxis.set_major_locator(MultipleLocator(5.893*1.8099*0.5*10))
        self.ax.xaxis.set_major_formatter(FormatStrFormatter('%d'))
        self.ax.yaxis.set_major_formatter(FormatStrFormatter('%1.1f'))
        self.ax.xaxis.set_minor_locator(MultipleLocator(0.5))
        self.ax.yaxis.set_minor_locator(MultipleLocator(5.893*1.8099*0.5))
        self.ax.set_xlim(0,30.0)  #设置坐标取值范围
        self.ax.set_ylim(y_range)
        legend = self.ax.legend(loc="upper right",prop=font_set)
        legend_f = legend.get_frame()
        legend_f.set_facecolor("white")
        
        
        
        
    def set_color_and_mark(self,color = list , mark = "." ):
        self.color = color
        self.markerstyle = mark
        
    def get_xy(self):
        index_dict = collections.defaultdict(list)
        crisscross_list = set([i for i in range(len(self.tb))])
        for i in range(1,self.type_num+1):
            for index,ratio in enumerate(self.type_dict[i]):
                if ratio >= 0.6321:
                    index_dict[i].append(index)
                    crisscross_list -= {index}
                    
        index_dict[0] = list(crisscross_list)
        for i in index_dict:
            yield (i,index_dict[i])
    
    def generate_image(self,save_path):
        
        for type_index,index_list in self.get_xy():
            if type_index == 0 and index_list:
                self.ax.plot(self.tb[index_list],self.E[index_list], "*", color=self.color[type_index], markersize = 8,label="交错类型")
                plt.legend(prop = font_set,loc = self.sign_loc)
            elif type_index != 0 :
                self.ax.plot(self.tb[index_list],self.E[index_list], self.markerstyle, color=self.color[type_index], markersize = 8,label="类型"+str(type_index))
                plt.legend(prop = font_set,loc = self.sign_loc)
        plt.savefig(save_path, dpi = 600,bbox_inches='tight')
        
        
class ExecuteCsv:
    def __init__(self,catalogue,title):

        a = re.compile(r'\d\.\d+\.\d\.\d+_?')
        self.file_dict = collections.defaultdict(list)
        self.title = title.replace("_"," ")
        
        def enumeration(dir1):
            for root,dirs,files in os.walk(dir1):        
                for file in files:   
                    if a.search(file):
                        csv_type = file.split("_")[-1].replace(".csv","") if "_" in file else file.split(".")[0]
                        csv_child_area = file.split(".")[1]
                        self.file_dict[(csv_type,csv_child_area)].append(os.path.join(root,file))
                for dir in dirs:            
                    enumeration(dir)   
                    
        enumeration(catalogue)
    
    def generate_csv(self,pathlist,save_path):
        dataMat = []
        for index,file in enumerate(pathlist):
            df = pd.read_csv(file)
            for i in range(len(df)):
                x = int((float(df.tb[i])/0.5 + 1))
                y = int((float(df.E_sum[i])/(5.893*1.8099*0.5) + 1))
                x_ret = x*0.5
                y_ret = y*(5.893*1.8099*0.5)

                retlist = [x_ret,y_ret] + [0]*len(pathlist)
                retlist[index+2] = 1
                dataMat.append(retlist)

        df = pd.DataFrame(dataMat)
        groupedMat = df.groupby([df[0],df[1]]).sum()
        groupedMat.to_csv(save_path) 

    def start_execute(self,save_path):
        for i in self.file_dict:
            true_title = self.got_true_title(i)
            true_path = save_path+f"/type={i[0]},district={i[1]}.csv"
            self.generate_csv(self.file_dict[i],true_path)
            latticeImage = LatticeImage(true_path,true_title)
            latticeImage.generate_image(save_path+f"/type={i[0]},district={i[1]}.png")
            

    def got_true_title(self,i = ("","")):
        tempr = ['温度交错带','高山极地带','苔原带','寒温带','中温带','暖温带','亚热带','热带']
        if len(i[0]) == 1:
            return self.title + " " + tempr[int(i[0])] + " 地域" + i[1]
        else:
            return self.title + " " + "-".join(tempr[int(j)] for j in i[0].split(".")) +"交错带" + " 地域" + i[1]
        
def test(csv_path,save_path):
    latticeImage = LatticeImage(csv_path,"test")
    latticeImage.generate_image(save_path)  
          
title = "387b_冷蒿、丛生矮禾草草原"
e = ExecuteCsv(r"E:\tutu\387b_冷蒿、丛生矮禾草草原",title) #做好的群落文件夹
e.start_execute(r"E:\tutu\cut") 
#test(r"E:\tutu\cut\5.0.1.1.csv",r"E:\tutu\cut\test.png")